from langchain_community.llms import Ollama
from config.index import Config

class LLMCheck:
    def __init__(self):
        self.cfg = Config()
        self.host = self.cfg.get('llm_check.host')
        self.model = self.cfg.get('llm_check.model')
        
        # 获取关键词列表
        self.keywords = self.cfg.get('keyword_search', [])
        
        # 初始化 phi3 模型，不使用 streaming 输出
        self.llm = Ollama(
            base_url=self.host,
            model=self.model,
            temperature=0.5  # 使用较低的温度以获得更确定的答案
        )
    
    def _check_keywords(self, query: str) -> bool:
        """
        检查查询中是否包含预定义的关键词
        """
        # 将查询转换为小写以进行不区分大小写的匹配
        query_lower = query.lower()
        
        # 检查每个关键词
        for keyword in self.keywords:
            if keyword.lower() in query_lower:
                print(f"关键词匹配成功: {keyword}")
                return True
        
        return False
        
    def needs_retrieval(self, query: str) -> bool:
        """
        判断是否需要调用知识库检索来回答问题
        1. 首先检查是否包含关键词
        2. 如果没有关键词匹配，则使用LLM进行判断
        """
        # 首先进行关键词检查
        if self._check_keywords(query):
            print("检测到关键词匹配，直接使用知识库检索")
            return True
            
        # 如果没有匹配到关键词，使用LLM进行判断
        prompt = f"""[INST]你是一个专业程序开发的判断器。
        规则：
        1. 只能回答 yes 或 no
        2. 不要解释原因
        3. 不要添加任何其他文字

        系统：判断以下问题是否需要开发技术文档知识库解答？
        
        问题："{query}"

        回答(只能说yes或no)：[/INST]"""
        
        try:
            response = self.llm.invoke(prompt).strip().lower()
            # 调试信息，仅显示实际收到的响应
            print(f"LLM原始响应: {response}")
            # 只保留第一个词，并检查是否为 yes
            first_word = response.split()[0] if response else ""
            result = first_word == 'yes'
            print(f"LLM判断结果: {'需要' if result else '不需要'}知识库检索")
            return result
        except Exception as e:
            print(f"LLM检查过程中发生错误: {e}")
            # 发生错误时默认使用知识库
            return True
